package core

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/base64"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"math/big"
	"mytool/base"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"
	"time"
)

func transfer(destination io.WriteCloser, source io.ReadCloser) {
	defer destination.Close()
	defer source.Close()
	io.Copy(destination, source)
}

func copyHeader(dst, src http.Header) {
	for k, vv := range src {
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

func handleHttps(w http.ResponseWriter, r *http.Request) {
	destConn, err := net.DialTimeout("tcp", r.Host, 60*time.Second)
	if err != nil {
		http.Error(w, err.Error(), http.StatusServiceUnavailable)
		return
	}
	w.WriteHeader(http.StatusOK)

	hijacker, ok := w.(http.Hijacker)
	if !ok {
		http.Error(w, "Hijacking not supported", http.StatusInternalServerError)
		return
	}

	clientConn, _, err := hijacker.Hijack()
	if err != nil {
		http.Error(w, err.Error(), http.StatusServiceUnavailable)
	}
	go transfer(destConn, clientConn)
	go transfer(clientConn, destConn)

}

func handleHttp(w http.ResponseWriter, r *http.Request) {
	resp, err := http.DefaultTransport.RoundTrip(r)
	if err != nil {
		http.Error(w, err.Error(), http.StatusServiceUnavailable)
		return
	}
	defer resp.Body.Close()

	copyHeader(w.Header(), resp.Header)
	w.WriteHeader(resp.StatusCode)
	io.Copy(w, resp.Body)

}

func parseProxyBasicAuth(r *http.Request) (username, password string, ok bool) {
	auth := r.Header.Get("Proxy-Authorization")
	if auth == "" {
		return "", "", false
	}

	split := strings.Split(auth, "Basic ")
	if len(split) != 2 {
		return "", "", false
	}

	c, err := base64.StdEncoding.DecodeString(split[1])
	if err != nil {
		return "", "", false
	}

	cs := string(c)
	username, password, ok = strings.Cut(cs, ":")
	if !ok {
		return "", "", false
	}
	return username, password, true
}

func checkIsBasicAuthOk(w http.ResponseWriter, r *http.Request, username, passwd string) bool {
	// 配置的空用户密码直接放行
	if username == "" && passwd == "" {
		return true
	}

	// 检查http basic认证信息
	u, p, ok := parseProxyBasicAuth(r)

	// 向客户端弹出认证
	if !ok {
		w.Header().Set("proxy-Authenticate", `Basic realm="restricted", charset="UTF-8"`)
		http.Error(w, "Unauthorized", http.StatusProxyAuthRequired)
		return false
	}

	return username == u && passwd == p
}

func doProxy(w http.ResponseWriter, r *http.Request) {
	// fmt.Println(r.RemoteAddr, "->", r.Host)
	if r.Method == http.MethodConnect {
		handleHttps(w, r)
	} else {
		handleHttp(w, r)
	}
}

func generateKeyPair() (rawCert, rawKey []byte, err error) {
	// Create private key and self-signed certificate
	// Adapted from https://golang.org/src/crypto/tls/generate_cert.go

	priv, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return
	}
	validFor := time.Hour * 24 * 365 * 10 // ten years
	notBefore := time.Now()
	notAfter := notBefore.Add(validFor)
	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	template := x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			Organization: []string{"Mytool"},
		},
		NotBefore: notBefore,
		NotAfter:  notAfter,

		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
	}
	derBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, &priv.PublicKey, priv)
	if err != nil {
		return
	}

	rawCert = pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: derBytes})
	rawKey = pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(priv)})

	return
}

func genCertificate() (cert tls.Certificate, err error) {
	rawCert, rawKey, err := generateKeyPair()
	if err != nil {
		return
	}
	return tls.X509KeyPair(rawCert, rawKey)

}

// HttpProxy 启动一个https代理 整理自：https://github.com/Zartenc/httpsproxy
func HttpProxy(args []string) (err error) {
	address := ":808"
	username := ""
	passwd := ""
	if len(args) >= 2 {
		address = args[1]
	}

	if len(args) == 3 {
		arr := strings.Split(args[2], ":")
		if len(arr) != 2 {
			return errors.New("账号密码格式不对，eg. proxy 808 admin:123456")
		}
		username = arr[0]
		passwd = arr[1]
	}

	cert, err := genCertificate()
	if err != nil {
		return
	}

	server := &http.Server{
		Addr:      address,
		TLSConfig: &tls.Config{Certificates: []tls.Certificate{cert}},
		Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			// fmt.Println(r, 666)
			fmt.Println(r.RequestURI)
			if strings.Contains(r.URL.Host, "hrfw.zhihui56.com") {

			}
			if checkIsBasicAuthOk(w, r, username, passwd) {
				doProxy(w, r)
			}
		}),
	}

	return server.ListenAndServe()
}

func HttpForward(args []string) (err error) {
	if len(args) != 3 {
		return errors.New("参数错误")
	}
	sourcePort := args[1]
	dstUrl := args[2]
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		u, _ := url.Parse(dstUrl)
		proxy := httputil.NewSingleHostReverseProxy(u)
		proxy.ServeHTTP(w, r)
	})
	err = http.ListenAndServe(":"+sourcePort, nil)
	return
}

func HttpProxyTools() []*base.Option {
	return []*base.Option{
		&base.Option{"http_proxy", HttpProxy, `http_proxy [ip:port] [admin:admin] http代理上网 eg. http_proxy 192.168.1.1:8080 admin:admin`},
		&base.Option{"http_forward", HttpForward, `http_forward <port> <url> http反向代理 eg. http_forward 8080 http://192.168.1.10:8080`},
	}
}
